import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt

# 自定义数据集类
class SinWaveDataset(Dataset):
    def __init__(self, num_samples, sequence_length, num_classes):
        self.num_samples = num_samples
        self.sequence_length = sequence_length
        self.num_classes = num_classes
        self.data, self.labels = self.generate_data()

    def generate_data(self):
        data = []
        labels = []
        frequencies = np.linspace(1, 10, self.num_classes)  # 不同频率

        for _ in range(self.num_samples):
            tmax = np.random.uniform(np.pi, 4 * np.pi)
            t = np.linspace(0, tmax, self.sequence_length)  # 时间点
            freq = np.random.choice(frequencies)  # 随机选择一个频率
            amplitude = np.random.uniform(0.5, 2.0)  # 随机振幅
            phase = np.random.uniform(0, 2 * np.pi)  # 随机相位
            p = amplitude * np.sin(freq * t + phase)  # 生成压力数值
            data.append(np.stack([t, p], axis=1))  # 每个样本是 (sequence_length, 2)
            labels.append(np.where(frequencies == freq)[0][0])  # 频率对应的类别

        data = np.array(data, dtype=np.float32)
        labels = np.array(labels, dtype=np.int64)
        return data, labels

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 定义模型
class TimeSeriesFeatureExtractor(nn.Module):
    def __init__(self, input_dim, embed_dim, num_heads, num_layers, num_classes):
        super(TimeSeriesFeatureExtractor, self).__init__()
        self.embedding = nn.Linear(input_dim, embed_dim)
        self.positional_encoding = nn.Parameter(torch.zeros(1, 50, embed_dim))
        #self.encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads)
        self.transformer_encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads)
                                                         , num_layers=num_layers)
        self.pooling = nn.AdaptiveAvgPool1d(1)
        self.fc = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        # x: (batch_size, sequence_length, input_dim)
        x = self.embedding(x)  # (batch_size, sequence_length, embed_dim)
        x = x + self.positional_encoding  # Add positional encoding
        x = x.permute(1, 0, 2)  # (sequence_length, batch_size, embed_dim)
        x = self.transformer_encoder(x)  # (sequence_length, batch_size, embed_dim)
        x = x.permute(1, 2, 0)  # (batch_size, embed_dim, sequence_length)
        x = self.pooling(x)  # (batch_size, embed_dim, 1)
        x = x.squeeze(-1)  # (batch_size, embed_dim)
        x = self.fc(x)  # (batch_size, num_classes)
        return x

# 训练函数
def train(model, dataloader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for batch_idx, (data, labels) in enumerate(dataloader):
        data, labels = data.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    avg_loss = total_loss / len(dataloader)
    accuracy = 100 * correct / total
    return avg_loss, accuracy

# 评估函数
def evaluate(model, dataloader, criterion, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0

    with torch.no_grad():
        for data, labels in dataloader:
            data, labels = data.to(device), labels.to(device)
            outputs = model(data)
            loss = criterion(outputs, labels)

            total_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    avg_loss = total_loss / len(dataloader)
    accuracy = 100 * correct / total
    return avg_loss, accuracy

# 主函数
def main():
    # 超参数
    input_dim = 2  # 每个时间步是 2 维向量 (时刻, 压力数值)
    embed_dim = 64  # 嵌入维度
    num_heads = 8  # 注意力头数
    num_layers = 2  # Transformer 编码器层数
    num_classes = 10  # 分类类别数
    batch_size = 256
    num_epochs = 400
    learning_rate = 0.001
    num_samples = 5000  # 训练样本数
    sequence_length = 50  # 时间序列长度

    # 数据集和数据加载器
    train_dataset = SinWaveDataset(num_samples, sequence_length, num_classes)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    test_dataset = SinWaveDataset(200, sequence_length, num_classes)
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # 设备设置
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 初始化模型
    model = TimeSeriesFeatureExtractor(input_dim, embed_dim, num_heads, num_layers, num_classes).to(device)

    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 训练循环
    for epoch in range(num_epochs):
        train_loss, train_acc = train(model, train_dataloader, criterion, optimizer, device)
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {train_loss:.4f}, Accuracy: {train_acc:.2f}%")

    # 评估模型
    test_loss, test_acc = evaluate(model, test_dataloader, criterion, device)
    print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.2f}%")

    # 保存模型
    torch.save(model.state_dict(), f"PlayGround/Attension2.pth")

if __name__ == "__main__":
    # 可视化一个样本
    dataset = SinWaveDataset(num_samples=5, sequence_length=50, num_classes=10)
    plt.figure(figsize=(10, 6))
    for i in range(5):
        plt.subplot(2, 3, i+1)
        t = dataset.data[i][:, 0]  # 时刻
        p = dataset.data[i][:, 1]  # 压力数值
        plt.plot(t, p)
        plt.title(f"Frequency Class: {dataset.labels[i]}")
    plt.tight_layout()
    plt.show()
    main()